#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
民宿价格预测模型
使用LightGBM进行民宿价格回归预测
"""

import time
import lightgbm as lgb
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from sklearn import metrics
from sklearn.model_selection import KFold
from sklearn.preprocessing import LabelEncoder
import joblib

class 民宿价格预测模型:
    """民宿价格预测模型类"""
    
    def __init__(self):
        self.model = None
        self.label_encoders = {}
        self.features = []
        
    def 数据预处理(self, 训练集路径='训练集.csv', 测试集路径='测试集.csv'):
        """数据预处理和特征工程"""
        print("=== 开始数据预处理 ===")
        
        # 读取数据
        self.train = pd.read_csv(训练集路径)
        self.test = pd.read_csv(测试集路径)
        
        print(f"训练集形状: {self.train.shape}, 测试集形状: {self.test.shape}")
        
        # 便利设施数量特征
        self.train['便利设施数量'] = self.train['便利设施'].apply(
            lambda x: len(str(x).lstrip('{').rstrip('}').split(',')) if pd.notna(x) else 0
        )
        self.test['便利设施数量'] = self.test['便利设施'].apply(
            lambda x: len(str(x).lstrip('{').rstrip('}').split(',')) if pd.notna(x) else 0
        )
        
        # 处理分类特征 - 标签编码
        categorical_features = self.train.select_dtypes(include=['object']).columns.tolist()
        categorical_features = [col for col in categorical_features if col not in ['数据ID', '便利设施']]
        
        for feature in categorical_features:
            le = LabelEncoder()
            # 合并训练集和测试集数据来拟合编码器，处理特殊字符
            combined = pd.concat([self.train[feature], self.test[feature]], axis=0)
            
            # 清理特殊字符和表情符号
            combined_clean = combined.astype(str).apply(
                lambda x: ''.join([c for c in x if ord(c) < 256 and c.isprintable()])
            )
            
            le.fit(combined_clean)
            self.train[feature] = le.transform(self.train[feature].astype(str).apply(
                lambda x: ''.join([c for c in str(x) if ord(c) < 256 and c.isprintable()])
            ))
            self.test[feature] = le.transform(self.test[feature].astype(str).apply(
                lambda x: ''.join([c for c in str(x) if ord(c) < 256 and c.isprintable()])
            ))
            self.label_encoders[feature] = le
        
        # 处理缺失值
        for df in [self.train, self.test]:
            for col in df.columns:
                if df[col].isnull().any():
                    if df[col].dtype == 'object':
                        df[col].fillna('Unknown', inplace=True)
                    else:
                        df[col].fillna(df[col].median(), inplace=True)
        
        print("数据预处理完成")
        
    def 特征选择(self):
        """选择用于训练的特征"""
        no_features = ['数据ID', '价格', '便利设施']
        self.features = [col for col in self.train.columns if col not in no_features]
        print(f"使用的特征: {self.features}")
        
        self.X = self.train[self.features]
        self.y = self.train['价格']
        self.X_test = self.test[self.features]
    
    def 训练模型(self, n_fold=5):
        """训练LightGBM模型"""
        print("=== 开始训练模型 ===")
        
        # 模型参数
        params = {
            'learning_rate': 0.1,
            'boosting_type': 'gbdt',
            'objective': 'regression',
            'metric': 'mae',
            'feature_fraction': 0.6,
            'bagging_fraction': 0.8,
            'bagging_freq': 5,
            'num_leaves': 1000,
            'verbose': -1,
            'max_depth': -1,
            'seed': 2019,
        }
        
        folds = KFold(n_splits=n_fold, shuffle=True, random_state=1314)
        oof = np.zeros(len(self.X))
        self.prediction = np.zeros(len(self.X_test))
        
        for fold_n, (train_index, valid_index) in enumerate(folds.split(self.X)):
            print(f"训练折数 {fold_n + 1}/{n_fold}")
            
            X_train, X_valid = self.X.iloc[train_index], self.X.iloc[valid_index]
            y_train, y_valid = self.y.iloc[train_index], self.y.iloc[valid_index]
            
            model = lgb.LGBMRegressor(**params, n_estimators=50000, n_jobs=-1, early_stopping_rounds=200)
            model.fit(
                X_train, y_train,
                eval_set=[(X_valid, y_valid)],
                eval_metric='rmse'
            )
            
            y_pred_valid = model.predict(X_valid)
            y_pred = model.predict(self.X_test, num_iteration=model.best_iteration_)
            
            oof[valid_index] = y_pred_valid.reshape(-1, )
            self.prediction += y_pred
            
            if fold_n == 0:
                self.model = model
        
        self.prediction /= n_fold
        print("模型训练完成")
    
    def 评估模型(self):
        """评估模型性能"""
        print("\n=== 模型评估 ===")
        
        # 计算训练集上的评估指标
        train_predictions = np.zeros(len(self.X))
        params = {
            'learning_rate': 0.1,
            'boosting_type': 'gbdt',
            'objective': 'regression',
            'metric': 'mae',
            'feature_fraction': 0.6,
            'bagging_fraction': 0.8,
            'bagging_freq': 5,
            'num_leaves': 1000,
            'verbose': -1,
            'max_depth': -1,
            'seed': 2019,
        }
        
        folds = KFold(n_splits=5, shuffle=True, random_state=1314)
        
        for fold_n, (train_index, _) in enumerate(folds.split(self.X)):
            X_train_fold = self.X.iloc[train_index]
            model_fold = lgb.LGBMRegressor(**params, n_estimators=self.model.best_iteration_)
            model_fold.fit(X_train_fold, self.y.iloc[train_index])
            train_predictions += model_fold.predict(self.X)
        
        train_predictions /= 5
        
        # 计算评估指标
        mae = metrics.mean_absolute_error(self.y, train_predictions)
        mse = metrics.mean_squared_error(self.y, train_predictions)
        rmse = np.sqrt(mse)
        r2 = metrics.r2_score(self.y, train_predictions)
        
        print(f"平均绝对误差 (MAE): {mae:.2f}")
        print(f"均方根误差 (RMSE): {rmse:.2f}")
        print(f"决定系数 (R²): {r2:.4f}")
    
    def 保存结果(self):
        """保存预测结果和模型"""
        print("\n=== 保存结果 ===")
        
        # 保存模型
        model_filename = '民宿价格预测模型.pkl'
        joblib.dump(self.model, model_filename)
        print(f"模型已保存到 {model_filename}")
        
        # 保存预测结果
        submission = pd.DataFrame({
            '数据ID': self.test['数据ID'],
            '预测价格': self.prediction
        })
        submission.to_csv('预测结果.csv', index=False, encoding='utf-8')
        print("预测结果已保存到 '预测结果.csv'")
        
        # 特征重要性分析
        feature_importance = pd.DataFrame({
            'feature': self.features,
            'importance': self.model.feature_importances_
        })
        feature_importance = feature_importance.sort_values('importance', ascending=False)
        
        plt.figure(figsize=(12, 8))
        plt.barh(feature_importance['feature'][:10], feature_importance['importance'][:10])
        plt.xlabel('特征重要性')
        plt.title('Top 10 特征重要性')
        plt.gca().invert_yaxis()
        plt.savefig('特征重要性.png', dpi=300, bbox_inches='tight')
        plt.close()
        print("特征重要性图表已保存到 '特征重要性.png'")
        
        # 预测结果统计
        prediction_stats = pd.DataFrame({
            '统计量': ['最小值', '最大值', '平均值', '中位数', '标准差'],
            '值': [
                self.prediction.min(),
                self.prediction.max(),
                self.prediction.mean(),
                np.median(self.prediction),
                self.prediction.std()
            ]
        })
        
        print("\n预测结果统计:")
        print(prediction_stats.to_string(index=False))
    
    def 运行完整流程(self):
        """运行完整的预测流程"""
        import sys
        import io
        
        # 设置标准输出为UTF-8
        sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
        
        start_time = time.time()
        print(f"开始时间: {time.strftime('%Y-%m-%d %H:%M:%S')}")
        
        try:
            print("\n步骤 1/5: 数据预处理")
            self.数据预处理()
            
            print("\n步骤 2/5: 特征选择") 
            self.特征选择()
            
            print("\n步骤 3/5: 训练模型 (可能需要几分钟...)")
            self.训练模型()
            
            print("\n步骤 4/5: 评估模型")
            self.评估模型()
            
            print("\n步骤 5/5: 保存结果")
            self.保存结果()
            
            end_time = time.time()
            print(f"\n✅ === 流程完成 ===")
            print(f"总运行时间: {end_time - start_time:.2f} 秒")
            print(f"完成时间: {time.strftime('%Y-%m-%d %H:%M:%S')}")
            
        except Exception as e:
            end_time = time.time()
            print(f"\n❌ === 流程失败 ===")
            print(f"错误信息: {e}")
            print(f"运行时间: {end_time - start_time:.2f} 秒")
            import traceback
            traceback.print_exc()

def main():
    """主函数"""
    print("=" * 50)
    print("民宿价格预测模型")
    print("=" * 50)
    
    # 创建模型实例
    模型 = 民宿价格预测模型()
    
    try:
        # 运行完整流程
        模型.运行完整流程()
        
        print("\n使用方法:")
        print("1. 将训练数据放入 '训练集.csv'")
        print("2. 将测试数据放入 '测试集.csv'") 
        print("3. 运行此脚本即可获得预测结果")
        
    except Exception as e:
        print(f"运行出错: {e}")
        print("请检查数据文件是否存在且格式正确")

if __name__ == "__main__":
    main()